package com.xiaoxin.experience.system.rate;

import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

/**
 * @author xiaoxin
 */
public class RateLimiter {

    public RateLimiter(long rate, long interval)
    {
        this.capacity = rate;
        this.rate = rate;
        this.interval = interval;
        this.tokenNum.set(rate);
    }

    /**
     * 桶的容量
     */
    private long capacity;

    /**
     * 每个周期放入令牌的数量
     */
    private long rate;

    /**
     * 令牌生成周期
     */
    private long interval;

    /**
     * 上次放置令牌的时间
     */
    private long lastTime = System.currentTimeMillis();

    /**
     * 桶中令牌的余量
     */
    private final AtomicLong tokenNum = new AtomicLong();

    /**
     * true 代表放行，请求可已通过
     * false 代表限制，不让请求通过
     */
    public synchronized boolean tryAcquire()
    {
        //更新桶中剩余令牌的数量
        long now = System.currentTimeMillis();
        tokenNum.addAndGet((now - lastTime) / interval * rate);
        tokenNum.set(Math.min(capacity, tokenNum.get()));
        //更新时间
        lastTime += (now - lastTime) / interval * interval;
        //桶中还有令牌就放行
        if (tokenNum.get() > 0)
        {
            tokenNum.decrementAndGet();
            return true;
        }
        else
        {
            return false;
        }
    }

    public void setRate(long rate, long interval, TimeUnit timeUnit)
    {
        this.rate = rate;
        this.tokenNum.set(rate);
        this.interval = timeUnit.toMillis(interval);
        this.capacity = rate;
    }

    public void setCapacity(long capacity)
    {
        this.capacity = capacity;
    }
}
